from nl2sql.datasource.base import BaseDatasource, Table, Column
import sqlite3
import pandas as pd
from typing import List, Dict, Optional, Any, Tuple
from nl2sql.similarity.cosine_similarity import CosineSimilarity
from nl2sql.model.embedding_model import EmbeddingModel


class SqliteDataSource(BaseDatasource):
    """支持加载SQLite数据库中所有表的数据源类"""

    def __init__(
            self,
            db_path: str,
            name: str = None,
            embedding_model: Optional[EmbeddingModel] = None,
            batch_size: int = 10
    ):
        """
        初始化SQLite数据源

        Args:
            db_path (str): SQLite数据库文件路径
            embedding_model (Optional[EmbeddingModel]): 嵌入模型
            batch_size (int): 批处理大小
        """

        super().__init__(
            embedding_model=embedding_model,
            batch_size=batch_size
        )
        self.name = name
        self.db_path = db_path
        self.connection: Optional[sqlite3.Connection] = None
        self.tables_data: Dict[str, pd.DataFrame] = {}  # 存储所有表的DataFrame
        self.table_names: List[str] = []  # 存储所有表名
        self.load_data()

    def load_data(self) -> None:
        """加载SQLite数据库中所有表的数据和结构信息"""
        if not self.db_path:
            raise ValueError("SQLite数据库路径未设置")

        try:
            # 建立数据库连接
            self.connection = sqlite3.connect(self.db_path)
            print(f"已连接到SQLite数据库: {self.db_path}")

            # 获取所有表名
            self.table_names = self._get_all_table_names()
            if not self.table_names:
                raise ValueError("数据库中没有找到表")

            # 加载所有表的数据和结构
            self.tables = []

            for table_name in self.table_names:
                table = Table()
                table.name = table_name
                # 加载列名
                columns: List[str] = self._get_column_names(table_name=table_name)
                for column_name in columns:
                    column = Column()
                    column.name = column_name
                    table.columns.append(column)
                self.tables.append(table)

            print(f"成功加载{len(self.table_names)}个表的数据和结构")

        except sqlite3.Error as e:
            print(f"加载SQLite数据失败: {e}")
            raise
        except Exception as e:
            print(f"加载过程中发生错误: {e}")
            raise

    def _get_column_names(self, table_name: str) -> List[str]:
        """
        根据表名获取列名
        """
        try:
            cursor = self.connection.cursor()
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = cursor.fetchall()
            column_names = [column[1] for column in columns]  # 提取列名
            print(f"表 {table_name} 的列名: {column_names}")
            return column_names
        except sqlite3.Error as e:
            print(f"获取表 {table_name} 的列名失败: {e}")
            return []

    def _get_all_table_names(self) -> List[str]:
        """获取数据库中所有表名"""
        try:
            cursor = self.connection.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';")
            table_names = [row[0] for row in cursor.fetchall()]
            print(f"找到{len(table_names)}个表: {table_names}")
            return table_names
        except sqlite3.Error as e:
            print(f"获取表名失败: {e}")
            return []

    def _load_single_table_data(self, table_name: str) -> None:
        """加载单个表的数据到DataFrame"""
        try:
            query = f"SELECT * FROM {table_name}"
            self.tables_data[table_name] = pd.read_sql_query(query, self.connection)
            print(f"已加载表: {table_name}，行数: {len(self.tables_data[table_name])}")
        except sqlite3.Error as e:
            print(f"加载表 {table_name} 失败: {e}")
            self.tables_data[table_name] = pd.DataFrame()  # 空DataFrame作为占位符

    def _get_table_structure(self, table_name: str) -> Optional[Table]:
        """获取单个表的结构信息"""
        try:
            cursor = self.connection.cursor()
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns_info = cursor.fetchall()

            # 获取表描述（SQLite默认不存储表注释，可扩展为外部元数据文件）
            table_desc = self._get_table_description(table_name) or f"Table {table_name}"

            table = Table(
                name=table_name,
                description=table_desc,
                columns=[]
            )

            for col_info in columns_info:
                col_name = col_info[1]
                # col_type = col_info[2]  # 可用于扩展类型信息

                # 获取列描述（SQLite默认不存储列注释，可扩展）
                col_desc = self._get_column_description(table_name, col_name) or f"Column {col_name}"

                # 获取枚举值
                enums = self._get_column_enums_from_data(table_name, col_name)

                column = Column(
                    name=col_name,
                    description=col_desc,
                    enums=self._filter_valid_enums(enums)  # 使用基类的过滤方法
                )
                table.columns.append(column)

            return table
        except sqlite3.Error as e:
            print(f"获取表 {table_name} 结构失败: {e}")
            return None

    def _get_table_description(self, table_name: str) -> Optional[str]:
        """获取表描述（可扩展为从外部元数据获取）"""
        # SQLite不原生支持表注释，这里返回None或默认值
        # 可扩展：从外部JSON文件或数据库读取描述
        return None

    def _get_column_description(self, table_name: str, column_name: str) -> Optional[str]:
        """获取列描述（可扩展为从外部元数据获取）"""
        # SQLite不原生支持列注释，这里返回None或默认值
        # 可扩展：从外部JSON文件或数据库读取描述
        return None

    def _get_column_enums_from_data(self, table_name: str, column_name: str) -> List[Any]:
        """从表数据中获取列的枚举值"""
        if table_name not in self.tables_data or column_name not in self.tables_data[table_name].columns:
            print(f"表 {table_name} 或列 {column_name} 不存在")
            return []

        data = self.tables_data[table_name]
        unique_values = data[column_name].dropna().unique()
        return list(unique_values)  # 返回所有唯一值，过滤在Column初始化时处理

    def get_table_data(self, table_name: str) -> Optional[pd.DataFrame]:
        """
        获取指定表的数据

        Args:
            table_name (str): 表名

        Returns:
            Optional[pd.DataFrame]: 表的DataFrame，不存在则返回None
        """
        if table_name not in self.tables_data:
            print(f"表 {table_name} 不存在")
            return None
        return self.tables_data[table_name].copy()

    def execute_custom_sql(self, sql: str) -> pd.DataFrame:
        """执行自定义SQL查询，支持跨表操作"""
        if not self.connection:
            self.connection = sqlite3.connect(self.db_path)

        try:
            return pd.read_sql_query(sql, self.connection)
        except sqlite3.Error as e:
            print(f"SQL执行失败: {e}")
            raise

    def get_all_table_names(self) -> List[str]:
        """获取所有已加载的表名"""
        return self.table_names.copy()

    def close_connection(self) -> None:
        """关闭数据库连接"""
        super().close_connection()  # 调用基类的关闭方法

    def get_data_info(self) -> Dict[str, Any]:
        """
        获取数据的核心信息，专为大模型生成SQL设计

        Returns:
            Dict[str, Any]: 包含SQL生成所需的数据信息
        """
        info = {
            "tables": []
        }

        for table in self.tables:
            table_info = {
                "name": table.name,
                "description": table.description,
                "columns": []
            }

            for column in table.columns:
                column_info = {
                    "name": column.name,
                    "description": column.description
                }
                table_info["columns"].append(column_info)

            info["tables"].append(table_info)

        return info


if __name__ == "__main__":
    sqlite_ds = SqliteDataSource(db_path='/Users/tiantiantian/bird_benchmark_dataset/dev/dev_20240627/dev_databases/california_schools/california_schools.sqlite',
                                 name='california_schools')

    info = sqlite_ds.get_data_info()
